import torch
import torch.nn as nn
from dual_input_net import DeepFusionNet, ProjectionHead, ClassificationHead

class FullModel(nn.Module):
    def __init__(self):
        super(FullModel, self).__init__()
        self.backbone = DeepFusionNet(in_channels=1, base_channels=32, feature_dim=128)
        self.proj_head = ProjectionHead(input_dim=128, projection_dim=128)
        self.cls_head = ClassificationHead(input_dim=128, hidden_dim=512, num_classes=3)

    def forward(self, b, h):
        features = self.backbone(b, h)       # [B, 128]
        projections = self.proj_head(features)  # [B, 128]
        logits = self.cls_head(features)     # [B, 3]
        return projections, logits
    
if __name__ == "__main__":
    model = FullModel()
    input = torch.randn(13, 128, 6, 7, 6)
    output = model(input)
    print(input.size(), output.size())